import torchvision
import numpy as np
from PIL import Image
import os
from dataset.utils_noise import noisify_cifar100_asymmetric
from dataset.utils_imbalance import get_imb_data
from dataset.imagenet32 import Imagenet32
import torch


class CIFAR10(torchvision.datasets.CIFAR10):
    nb_classes = 10

    def __init__(self, root='~/data', train=True, transform=None,
                 r_ood=0.2, r_id=0.2, r_imb=0.1, seed=0, asym=False,
                 mode=None, pred=None, probability=None, transform_st=None):

        super().__init__(root, train=train, transform=transform)
        if train is False:
            return
        if r_imb > 0.:
            self.data, self.targets = get_imb_data(self.data, self.targets,
                                                   self.nb_classes, r_imb, seed)
            print('built imbalanced dataset')
        np.random.seed(seed)
        if r_ood > 0.:
            ids_ood = [i for i in range(len(self.targets)) if np.random.random() < r_ood]
            imagenet32 = Imagenet32(root='~/data/imagenet32', train=True)
            img_ood = imagenet32.data[np.random.permutation(range(len(imagenet32)))[:len(ids_ood)]]
            self.ids_ood = ids_ood
            self.data[ids_ood] = img_ood
            print('Mixing in OOD noise')

        if r_id > 0.:
            ids_not_ood = [i for i in range(len(self.targets)) if i not in ids_ood]
            ids_id = [i for i in ids_not_ood if np.random.random() < (r_id / (1 - r_ood))]
            if asym:
                if self.nb_classes == 10:
                    transition = {0: 0, 2: 0, 4: 7, 7: 7, 1: 1, 9: 1, 3: 5, 5: 3, 6: 6, 8: 8}
                    for i, t in enumerate(self.targets):
                        if i in ids_id:
                            self.targets[i] = transition[t]
                else:
                    self.targets = noisify_cifar100_asymmetric(self.targets, r_id, seed)
                print('Mixing in ID asym noise')
            else:
                for i, t in enumerate(self.targets):
                    if i in ids_id:
                        self.targets[i] = int(np.random.random() * self.nb_classes)
                print('Mixing in ID noise')
            self.ids_id = ids_id
        self.mode = mode
        self.probability = probability
        self.pred = pred
        self.transform_st = transform_st
        if mode == 'labeled':
            self.filter_labeled(pred)
        elif mode == 'unlabeled':
            self.filter_unlabeled(pred)

    def filter_labeled(self, pred):
        pred_idx = pred.nonzero()[0]
        self.data = self.data[pred_idx]
        self.targets = np.array(self.targets)[pred_idx]

    def filter_unlabeled(self, pred):
        pred_idx = (1 - pred).nonzero()[0]
        self.data = self.data[pred_idx]
        self.targets = np.array(self.targets)[pred_idx]

    def __getitem__(self, index):
        assert self.transform is not None, 'Transform is None!'
        img, target = self.data[index], self.targets[index]

        img = Image.fromarray(img)
        target = torch.tensor(target).long()

        if self.train is False:
            return self.transform(img), target

        if self.mode is None:
            return self.transform(img), target, index
        img1 = self.transform(img)
        img2 = self.transform(img)
        if self.mode == 'labeled':
            prob = self.probability[index]
            return img1, img2, target, prob
        elif self.mode == 'unlabeled':
            # img1 = self.transform_st(img)
            img2 = self.transform_st(img)
            return img1, img2


class CIFAR100(CIFAR10):
    base_folder = "cifar-100-python"
    url = "https://www.cs.toronto.edu/~kriz/cifar-100-python.tar.gz"
    filename = "cifar-100-python.tar.gz"
    tgz_md5 = "eb9058c3a382ffc7106e4002c42a8d85"
    train_list = [
        ["train", "16019d7e3df5f24257cddd939b257f8d"],
    ]

    test_list = [
        ["test", "f0ef6b0ae62326f3e7ffdfab6717acfc"],
    ]
    meta = {
        "filename": "meta",
        "key": "fine_label_names",
        "md5": "7973b15100ade9c7d40fb424638fde48",
    }

    nb_classes = 100
